import os
from skimage import util
import cv2
from PIL import Image


def generate_patches_walk(patch_size=32, stride=14,
                          saved_path='/media/wjsun/delldisk/dell/wxm/Data/decsai/CM512/patches32val_dip'):
    train_pgms = []
    pgm_folder = '/media/wjsun/delldisk/dell/wxm/Data/decsai/CM512/val'
    flag_pgm = 0
    x_one = 0
    y_one = 0
    x_two = patch_size
    y_two = patch_size
    # nb_patches = 1000
    for dirpath, dirname, filname in os.walk(pgm_folder):
        for img in filname:
            patch_flag = 0
            flag_pgm += 1
            print 'nb.{0} images'.format(flag_pgm)
            train_pgms.append(img)
            img_path = dirpath + '/' + img
            img_obj = Image.open(img_path)
            for i in range(int((512 - patch_size) / stride + 1)):
                for j in range(int((512 - patch_size) / stride + 1)):
                    patch_flag += 1
                    box = (x_one + j * stride, y_one + i * stride, x_two + j * stride, y_two + i * stride)
                    img_patch = img_obj.crop(box)
                    img_patch.save(saved_path + '/' + str(patch_flag) + '_' + img)
                    if patch_flag % 200 == 0:
                        print 'nb.{0} patches'.format(patch_flag)


def load_train32C(Debug=False):
    src_folder = '/media/wjsun/delldisk/dell/wxm/Data/decsai/CM512/patches32train_dip'
    img_names = []
    imgs = []
    flag = 0
    for img in os.listdir(src_folder):
        img_path = src_folder + '/' + img
        img_names.append(img)
        img_arr = cv2.imread(img_path, 0)
        imgs.append(img_arr)
        flag += 1
        if Debug is True and flag >= 10000:
            break
    return imgs, img_names


def load_val32C(Debug=False):
    src_folder = '/media/wjsun/delldisk/dell/wxm/Data/decsai/CM512/patches32val_dip'
    img_names = []
    imgs = []
    flag = 0
    for img in os.listdir(src_folder):
        img_path = src_folder + '/' + img
        img_names.append(img)
        img_arr = cv2.imread(img_path, 0)
        imgs.append(img_arr)
        flag += 1
        if Debug is True and flag >= 5000:
            break
    return imgs, img_names


def gaussian_noise(X):
    for i in range(len(X)):
        X[i] = util.random_noise(X[i], 'gaussian', var=0.05)
    return X
